// +------------------------------------------------------------------------+
// | XRes - Mac OS X Resource Monitor                                       |
// +------------------------------------------------------------------------+
// | Copyright (c) 2007 Zack Bartel                                         |
// +------------------------------------------------------------------------+
// | This program is free software; you can redistribute it and/or          |
// | modify it under the terms of the GNU General Public License            | 
// | as published by the Free Software Foundation; either version 2         | 
// | of the License, or (at your option) any later version.                 |
// |                                                                        |
// | This program is distributed in the hope that it will be useful,        |
// | but WITHOUT ANY WARRANTY; without even the implied warranty of         |
// | MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the          |
// | GNU General Public License for more details.                           |
// |                                                                        |
// | You should have received a copy of the GNU General Public License      |
// | along with this program; if not, write to the Free Software            |
// | Foundation, Inc., 59 Temple Place - Suite 330,                         |
// | Boston, MA  02111-1307, USA.                                           |
// +------------------------------------------------------------------------+
// | Author: Zack Bartel <zack@bartel.com>                                  |
// +------------------------------------------------------------------------+ 
//
//  XRMemoryGraphDataSource.m
//  XRes
//
//  Created by Zack Bartel on 3/7/08.
//  Copyright 2008 Zack Bartel. All rights reserved.
//

#import <ScreenSaver/ScreenSaver.h>
#import "XRMemoryGraphDataSource.h"
#import "XRProcess.h"
#import "XRColorCell.h"



@implementation XRMemoryGraphDataSource

- (void) awakeFromNib
{
    selectedProcess = 0;
    reverseOrderBy = NO;
    monitorType = XRMonitorRam;
    
    processAttributes = XRProcessAllProcesses;
    processOrderBy    = XRProcessOrderByResidentSize;
    
    graphCPUData    =   [[NSMutableDictionary alloc] init];
    graphMemoryData =   [[NSMutableDictionary alloc] init];
    
    [stopButton setEnabled: NO];
    
    [self loadProcesses: nil];
    
    [chart setExplodeDistance: 20.0];
    [chart setDataSource: self];
    [chart refreshDisplay: self];
    
    [graph setDataSource: self];
    [graph refreshDisplay: self];
}

#pragma mark -
#pragma mark • UTILITY METHODS
- (void) loadProcesses: (NSTimer *) timer
{
    @synchronized (self)
    {
        if (processes)
        {
            [processes release];
            processes = nil;
        }
        
         processes = [[XRProcess allProcessesWithAttributes: processAttributes 
                                                  orderedBy: processOrderBy] 
                                                  retain];
    }
    
    if (reverseOrderBy)
    {
        NSMutableArray *newProcesses = nil;
        
        newProcesses = [[NSMutableArray arrayWithArray: [[processes reverseObjectEnumerator] allObjects]] retain];
        
        [processes release];
        
        processes = newProcesses;
    }
    
    [self selectProcess: nil];
    
     //Save current process state data
    [self saveCurrentStateGraphData];

    [chart refreshDisplay: self];
    
    [graph refreshDisplay: self];
    
    [table reloadData];
    
    if (selectedProcess > 0)
    {
        [table selectRowIndexes: [NSIndexSet indexSetWithIndex: [self indexOfProcess: selectedProcess]] byExtendingSelection: NO];
    }
    else
    {
        [table deselectAll: self];
    }
}

- (NSColor *) colorForProcess: (XRProcess *) process
{
    NSColor *color;
    
    NSNumber *pid = [NSNumber numberWithInt: [process pid]];
    
    if (colorMap == nil)
    {
        colorMap = [[NSMutableDictionary dictionary] retain];
    }
    
    color = [colorMap objectForKey: pid];
    
    if (color == nil)
    {
        color = [self getNextColor];
        [colorMap setObject: color forKey: pid];
    }
    
    return color;
}

- (int) indexOfProcess: (int) pid
{
    XRProcess *process = nil;
    int selectedProcessIndex = 0;
    
    @synchronized (self)
    {
        int i;
        int count = [processes count];
        
        for (i = 0; i < count; i++)
        {
            process = [processes objectAtIndex: i];
            
            if ([process pid] == selectedProcess)
            {
                selectedProcessIndex = i;
                break;
            }
        }
    }
    
    return selectedProcessIndex;
}

- (XRProcess *) processForPid: (int) pid
{
    XRProcess *process = nil;
    
    @synchronized (self)
    {
        int i;
        int count = [processes count];
        
        for (i = 0; i < count; i++)
        {
            process = [processes objectAtIndex: i];
            
            if ([process pid] == pid)
            {
                break;
            }
        }
    }
    
    return process;
}

- (void) setSelectedProcess: (int) value
{
    selectedProcess = value;
}

- (void) setMonitorType: (XRMonitorType) type
{
    monitorType = type;
    
    switch (monitorType)
    {
        case XRMonitorCpu:
            processOrderBy = XRProcessOrderByPercentCPUUsage;
            break;
        
        case XRMonitorRam:
            processOrderBy = XRProcessOrderByResidentSize;
            break;
            
        default:
            processOrderBy = XRProcessOrderByResidentSize;
            break;
    }
    
    reverseOrderBy = NO;
    
    [chart refreshDisplay: self];
    [graph refreshDisplay: self];
}

- (XRProcess *) selectedProcess
{
    return [self processForPid: selectedProcess];
}

- (NSColor *) getNextColor
{
    static int index = 0;
    
    NSColor *color;
    
    switch(index)
    {
        case 0:
            color = [NSColor redColor];
            break;
        case 1:
            color = [NSColor orangeColor];
            break;
        case 2:
            color = [NSColor yellowColor];
            break;
        case 3:
            color = [NSColor blueColor];
            break;
        case 4:
            color = [NSColor purpleColor];
            break;
        case 5:
            color = [NSColor brownColor];
            break;
        case 6:
            color = [NSColor cyanColor];
            break;
        case 7:
            color = [NSColor darkGrayColor];
            break;
        case 8:
            color = [NSColor grayColor];
            break;
        case 9:
            color = [NSColor greenColor];
            break;
        case 10:
            color = [NSColor lightGrayColor];
            break;
        case 11:
            color = [NSColor magentaColor];
            break;
        case 12:
            color = [NSColor blackColor];
            break;
        default:
            color = [NSColor colorWithCalibratedRed:(SSRandomFloatBetween( 0.0, 255.0 ) / 255.0)
                                green:(SSRandomFloatBetween( 0.0, 255.0 ) / 255.0)
                                blue:(SSRandomFloatBetween( 0.0, 255.0 ) / 255.0)
                                alpha:1.0];//(SSRandomFloatBetween( 0.0, 255.0 ) / 255.0)];
            break;
    }
    
    index++;
    
    return color;
}

- (void) saveCurrentStateGraphData
{
    NSNumber *pid = nil;
    NSMutableArray *cpuData = nil;
    NSMutableArray *memoryData = nil;
    NSPoint  point;
    
    @synchronized (self)
    {
        for (XRProcess *process in processes)
        {
            pid = [NSNumber numberWithInt: [process pid]];
            
            cpuData     = [graphCPUData objectForKey: pid];
            memoryData  = [graphMemoryData objectForKey: pid];
            
            if (cpuData == nil)
            {
                cpuData = [NSMutableArray array];
            }
            if (memoryData == nil)
            {
                memoryData = [NSMutableArray array];
            }
            
            if ([cpuData count] > 100)
            {
                [self popArray: cpuData];
            }
            if ([memoryData count] > 100)
            {
                [self popArray: memoryData];
            }
                                    
            point = NSMakePoint((float)[cpuData count], [process percentCPUUsage]);
            [cpuData addObject: NSStringFromPoint(point)];
            
            point = NSMakePoint((float)[memoryData count], [process percentResidentMemory]);
            [memoryData addObject: NSStringFromPoint(point)];
            
            [graphCPUData setObject: cpuData forKey: pid];
            [graphMemoryData setObject: memoryData forKey: pid];
        }
    }
}

- (void) popArray: (NSMutableArray *) data
{
    int i, count = [data count];
    
    for (i = 0; i < count - 1; i++)
    {
        NSPoint p = NSPointFromString([data objectAtIndex: i + 1]);
        [data replaceObjectAtIndex: i withObject: NSStringFromPoint(NSMakePoint( p.x - 1.0, p.y))];
    }
    
    [data removeLastObject];
}

//If process is nil then use the current selected PID
- (void) selectProcess: (XRProcess *) process
{
    if (process == nil)
    {
        process = [self processForPid: selectedProcess];
    }
    else if ([process pid] != selectedProcess)
    {
        selectedProcess = [process pid];
        [processLabel setStringValue: [process command]];
        [commandLabel setStringValue: [process path]];
        [commandLabelLabel setHidden: NO];
        
        [userLabel setStringValue: [process user]];
        [userLabelLabel setHidden: NO];
    
        [pidLabel setStringValue: [NSString stringWithFormat: @"%d", [process pid]]];
        [pidLabelLabel setHidden: NO];
        
        [processIcon setImage: [[NSWorkspace sharedWorkspace] iconForFile: [process path]]];
        [table selectRowIndexes: [NSIndexSet indexSetWithIndex: [self indexOfProcess: selectedProcess]] byExtendingSelection: NO];
        [stopButton setEnabled: YES];
    }
    
    if (selectedProcess == 0)
    {
        [processLabel setStringValue: @""];
        [commandLabel setStringValue: @""];
        [commandLabelLabel setHidden: YES];
        
        [userLabel setStringValue: @""];
        [userLabelLabel setHidden: YES];
    
        [pidLabel setStringValue: @""];
        [pidLabelLabel setHidden: YES];
        
        [cpuLabel setStringValue: @""];
        [cpuLabelLabel setHidden: YES];
    
        [ramLabel setStringValue: @""];
        [ramLabelLabel setHidden: YES];
        
        [processIcon setImage: nil];
        [table deselectAll: self];
        [stopButton setEnabled: NO];

    }
    else
    {
        [cpuLabel setStringValue: [NSString stringWithFormat: @"%.1f %%", [process percentCPUUsage]]];
        [cpuLabelLabel setHidden: NO];
    
        [ramLabel setIntValue: [process residentMemorySize]];
        [ramLabelLabel setHidden: NO];
    }
}

+ (NSString *) bytesToString: (long) numBytes
{
    double result = (double)numBytes;
    int level = 0;
    while((int)(result / 1024) > 0){
        result = (float)result / 1024;
        //keepGoing = YES;
        level++;
    }
    NSString *postFix;
    switch (level){
        case 0: //bytes (b)
            postFix = [[NSString alloc] initWithString:@"B"];
            break;
        case 1: //kilobytes (kB)
            postFix = [[NSString alloc] initWithString:@"K"];
            break;
        case 2: //megabytes (MB)
            postFix = [[NSString alloc] initWithString:@"MB"];
            break;
        case 3: //gigabytes (GB)
            postFix = [[NSString alloc] initWithString:@"GB"];
            break;
        case 4: //terabytes (TB)
            postFix = [[NSString alloc] initWithString:@"TB"];
            break;
        case 5: //petabytes (PB)
            postFix = [[NSString alloc] initWithString:@"PB"];
            break;
        case 6: //exabytes (EB)
            postFix = [[NSString alloc] initWithString:@"EB"];
            break;
        case 7: //zettabyte (ZB)
            postFix = [[NSString alloc] initWithString:@"ZB"];
            break;
        case 8: //yottabyte (YB)
            postFix = [[NSString alloc] initWithString:@"YB"];
            break;
        default:
            postFix = [[NSString alloc] initWithString:@"b"];
    }
    return [NSString stringWithFormat:@"%.0f%@", result, [postFix autorelease]];
}

- (void) runQuitSheet
{
    XRProcess *process = [self selectedProcess];
    NSBeginAlertSheet(@"Quit Process", @"Quit", @"Force Quit", @"Cancel", 
            [table window], self, 
            @selector(quitSheetClosed:returnCode:contextInfo:),
            NULL, NULL, [NSString stringWithFormat: @"Are you sure you want to quit the process \"%@\"?", [process command]]);
}

- (void) quitSheetClosed: (NSWindow *) sheet
         returnCode: (int) returnCode
        contextInfo: (void *) contextInfo
{
    int result = 0;
    
    switch (returnCode)
    {
        case NSAlertDefaultReturn:
            result = kill(selectedProcess, SIGTERM);
            break;
        case NSAlertAlternateReturn:
            result = kill(selectedProcess, SIGKILL);
            break;
        default:
            break;
    }
    
    if (result < 0)
    {
        //TODO
        NSLog(@"Error: Operation not permitted");
    }
    else
    {
        [self loadProcesses: nil];
    }
}

#pragma mark -
#pragma mark • MEMORY PIECHART DATASOURCE METHODS
- (unsigned int) numberOfSlicesInPieChartView: (SMPieChartView *) inPieChartView
{
    unsigned int count = 0;
    
    @synchronized (self)
    {
        count = [processes count];
    }
    
    return count;
}

- (double) pieChartView: (SMPieChartView *) inPieChartView dataForSliceIndex: (unsigned int) inSliceIndex
{
    double result = 0.0;
    XRProcess *process = nil;
    
    @synchronized (self)
    {
        process = [processes objectAtIndex: inSliceIndex];
        
        switch (monitorType)
        {
            case XRMonitorRam:
                result = [process percentResidentMemory];
                break;
                
            case XRMonitorCpu:
                result = [process percentCPUUsage];
                break;
                
            default:
                result = [process percentResidentMemory];
                break;
        }
    }
    
    return result;
}

- (unsigned int) numberOfExplodedPartsInPieChartView: (SMPieChartView *) inPieChartView
{
    return selectedProcess > 0 ? 1 : 0;
}

- (NSRange) pieChartView: (SMPieChartView *) inPieChartView rangeOfExplodedPartIndex: (unsigned int) index
{

    return NSMakeRange([self indexOfProcess: selectedProcess], 1);
}

- (NSDictionary *) pieChartView: (SMPieChartView *) inPieChartView attributesForSliceIndex: (unsigned int) inSliceIndex
{
    NSColor *color;
    
    @synchronized (self)
    {
        XRProcess *process = [processes objectAtIndex: inSliceIndex];
        
        color = [self colorForProcess: process];
    }
    
    return [ NSDictionary dictionaryWithObject:color forKey:NSBackgroundColorAttributeName ];
}

#pragma mark -
#pragma mark • MEMORY PIECHART DELEGATE METHODS
- (void) pieChartView: (SMPieChartView *) pieChartView didClickPoint: (NSPoint) point
{
    int slice = [pieChartView convertToSliceFromPoint: point fromView: pieChartView];
    
    @synchronized (self)
    {
        XRProcess *process = [processes objectAtIndex: slice];
        
        if (selectedProcess == [process pid])
        {
            selectedProcess = 0;
            [self selectProcess: nil];
        }
        else
        {
            [self selectProcess: process];
        }
    }
    
    [chart refreshDisplay: self];
}

#pragma mark -
#pragma mark • CPU GRAPH DATASOURCE METHODS
- (unsigned int) numberOfLinesInTwoDGraphView: (SM2DGraphView *) inGraphView
{
    unsigned int count = 0;
    
    @synchronized (self)
    {
        count = [processes count];
    }
    
    return count;

}

- (NSArray *) twoDGraphView: (SM2DGraphView *) inGraphView dataForLineIndex: (unsigned int) inLineIndex
{
    NSArray *result = nil;
    if (monitorType == XRMonitorCpu)
        result = [graphCPUData objectForKey: [NSNumber numberWithInt: [[processes objectAtIndex: inLineIndex] pid]]];
    else
        result = [graphMemoryData objectForKey: [NSNumber numberWithInt: [[processes objectAtIndex: inLineIndex] pid]]];
    
    return result;
}

- (double) twoDGraphView: (SM2DGraphView *) inGraphView maximumValueForLineIndex: (unsigned int) inLineIndex
            forAxis:(SM2DGraphAxisEnum)inAxis
{
    return 100.0;
}

- (double) twoDGraphView: (SM2DGraphView *) inGraphView minimumValueForLineIndex: (unsigned int) inLineIndex
            forAxis:(SM2DGraphAxisEnum)inAxis
{
    return 0.0;
}

- (NSDictionary *) twoDGraphView: (SM2DGraphView *) inGraphView attributesForLineIndex: (unsigned int) inLineIndex
{
    XRProcess       *process = nil;
    NSDictionary	*result = nil;
    
    @synchronized (self)
    {
        process = [processes objectAtIndex: inLineIndex];
        
        // Make the cosine line red and don't anti-alias it.
        result = [ NSDictionary dictionaryWithObjectsAndKeys:
                    [self colorForProcess: process], NSForegroundColorAttributeName,
                    //[ NSNumber numberWithBool:YES ], SM2DGraphDontAntialiasAttributeName,
                    nil ];
    }
    
    return result;
}

#pragma mark -
#pragma mark • PROCESS TABLEVIEW DATASOURCE METHODS
- (int) numberOfRowsInTableView: (NSTableView *) aTableView
{
    int result = 0;
    
    @synchronized (self)
    {
        result = [processes count];
    }
    
    return result;
}

- (id) tableView: (NSTableView *) tableView
    objectValueForTableColumn: (NSTableColumn *) tableColumn
    row: (int) row
{
    id result = nil;
    XRProcess *process = nil;
    NSString *identifier = [tableColumn identifier];
    
    @synchronized (self)
    {
        if (processes != nil)
        {
            process = [processes objectAtIndex: row];
            
            if (process != nil)
            {
                if ([identifier isEqualToString: @"processCommand"])
                {
                    result = [process command];
                }
                else if ([identifier isEqualToString: @"processColor"])
                {
                    result = [[[XRColorCell alloc] initWithColor: [self colorForProcess: process]] autorelease];
                    [tableColumn setDataCell: result];
                }
                else if ([identifier isEqualToString: @"percentCpu"])
                {
                    result = [NSString stringWithFormat: @"%0.01f", [process percentCPUUsage]];
                }
                else if ([identifier isEqualToString: @"percentMemory"])
                {
                    result = [NSString stringWithFormat: @"%0.01f", [process percentResidentMemory]];
                }
                else if ([identifier isEqualToString: @"virtualSize"])
                {
                    result = [XRMemoryGraphDataSource bytesToString: [process virtualSize]];
                }
                else if ([identifier isEqualToString: @"memorySize"])
                {
                    result = [XRMemoryGraphDataSource bytesToString: [process residentMemorySize]];
                }
                else if ([identifier isEqualToString: @"command"])
                {
                    result = [process path];
                }
                else if ([identifier isEqualToString: @"user"])
                {
                    result = [process user];
                }
            }
        }
    }
    
    return result;
}

- (void) tableView: (NSTableView *) tableView didClickTableColumn: (NSTableColumn *) tableColumn
{
    NSString *identifier = [tableColumn identifier];
    
    XRProcessOrderBy newProcessOrderBy = processOrderBy;
    
    if ([identifier isEqualToString: @"processCommand"])
    {
        newProcessOrderBy = XRProcessOrderByCommandName;
    }
    else if ([identifier isEqualToString: @"percentCpu"])
    {
        newProcessOrderBy = XRProcessOrderByPercentCPUUsage;
    }
    else if ([identifier isEqualToString: @"percentMemory"])
    {
        newProcessOrderBy = XRProcessOrderByResidentSize;
    }
    else if ([identifier isEqualToString: @"memorySize"])
    {
        newProcessOrderBy = XRProcessOrderByResidentSize;
    }
    else if ([identifier isEqualToString: @"virtualSize"])
    {
        newProcessOrderBy = XRProcessOrderByVirtualMemorySize;
    }
    else if ([identifier isEqualToString: @"user"])
    {
        newProcessOrderBy = XRProcessOrderByUser;
    }
    
    if (newProcessOrderBy == processOrderBy)
    {
        reverseOrderBy = !reverseOrderBy;
    }
    else
    {
        reverseOrderBy = NO;
    }
    
    processOrderBy = newProcessOrderBy;
    
    [self loadProcesses: nil];
}

#pragma mark -
#pragma mark • IBACTIONS
- (IBAction) processTableViewClicked: (id) sender
{
    @synchronized (self)
    {
        int selectedRow = [sender selectedRow];
        XRProcess *process = nil;
        
        if (selectedRow >= 0)
        {
            process = [processes objectAtIndex: [sender selectedRow]];
            [self selectProcess: process];
        }
    }
    
    [chart refreshDisplay: self];
}

- (IBAction) stop: (id) sender
{
   [self runQuitSheet];
}

- (IBAction) processFilterChanged: (id) sender
{
    XRProcessAttribute attribute; 
    int selected = [sender indexOfSelectedItem];
    
    switch (selected)
    {
        case 0:
            attribute = XRProcessAllProcesses;
            break;
        case 1:
            attribute = XRProcessUserProcessesOnly;
            break;
        case 2:
            attribute = XRProcessWindowedProcessesOnly;
            break;
        default:
            break;
    }
    
    if (attribute != processAttributes)
    {
        processAttributes = attribute;
        selectedProcess = 0;
        
        [self loadProcesses: nil];
    }
}

@end
